import gradio as gr
import numpy as np
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
from PIL import Image


def extract_color_palette(image):
    # 将图像转换为 numpy 数组
    img = np.array(image)

    # 将图像 reshape 为 (n_pixels, 3)，以适应 KMeans 输入格式
    reshaped_img = img.reshape((-1, 3))

    # 使用 KMeans 算法提取颜色调色板
    kmeans = KMeans(n_clusters=5)
    kmeans.fit(reshaped_img)

    # 获取聚类中心作为主要颜色
    colors = kmeans.cluster_centers_.astype(int)

    # 转换颜色为十六进制表示形式
    hex_colors = '\n'.join(['#%02x%02x%02x' %  
                            (color[0], color[1], color[2]) for color in colors])
    # 绘制颜色条
    palette = np.zeros((200, 1000, 3), dtype=np.uint8)
    for i, color in enumerate(colors):
        palette[:, i * 200:(i + 1) * 200] = color

    return palette, hex_colors


iface = gr.Interface(
    fn=extract_color_palette,
    inputs=gr.components.Image(),
    outputs=[
        gr.components.Image(type="numpy", show_label=False, width=1000),
        gr.components.Text(label="十六进制表示形式")
    ],
    title="KMeans调色板提取器",
    description="上传图像获取其调色板"
)

iface.launch()  # 在本地运行 Gradio 应用程序
